#! /usr/bin/python
# _*_ coding:UTF-8 _*_

import numpy as np
import matplotlib.pyplot as plt
import random


def k_means(k, data):
    (m, n) = data.shape
    c = np.zeros(m, int)
    center = np.zeros([k, n])
    # center = np.array([[0.403, 0.237], [0.343, 0.099], [0.532, 0.472]])
    # 在数据集中随机选择k个样本作为初始均值向量
    for i in range(0, k, 1):
        center[i] = data[random.randrange(0, m)]
    max_repater = 20  # 最大循环次数
    rep = 0

    while rep < max_repater:
        # print center
        # 计算欧式距离将样本划分到对应的簇
        for j in range(0, m, 1):
            temp = center - data[j].reshape(1, n)
            d_sum = np.sum(temp * temp, axis=1)  # 按行求和
            c[j] = np.where(d_sum == np.min(d_sum))[0][0]
        print c
        # 重新计算均值
        flag = 0
        for i in range(0, k, 1):
            # 第i簇样本个数
            center_i_temp = (1.0/np.sum(c == i)) * np.sum(data[np.where(c == i)], axis=0)
            if np.sum(center_i_temp == center[i]) != n :
                center[i] = center_i_temp
                flag += 0
            else:
                flag += 1
        if flag == k:
            break
        rep += 1
    return c, center

if __name__ == "__main__":
    print "开始"
    data = np.loadtxt("data.txt")
    # print data
    plt.figure(1)
    plt.plot(data[:, 0], data[:, 1], 'bo')
    plt.xlim(0.1, 0.9)
    plt.ylim(0, 0.8)
    plt.xlim()
    plt.xlabel(u"密度")
    plt.ylabel(u"含糖率")
    # plt.show()
    k = 3  # 类镞数
    (c, center) = k_means(k, data)
    plt.plot(data[np.where(c==0), 0], data[np.where(c==0), 1], 'bo')
    plt.plot(data[np.where(c==1), 0], data[np.where(c==1), 1], 'ro')
    plt.plot(data[np.where(c==2), 0], data[np.where(c==2), 1], 'yo')
    # print center
    plt.plot(center[0, 0], center[0, 1], 'b*', ms=10, label='center1')
    plt.plot(center[1, 0], center[1, 1], 'r*', ms=10, label='center2')
    plt.plot(center[2, 0], center[2, 1], 'y*', ms=10, label='center3')
    plt.legend()
    plt.show()



